package com.gitee.xmhzzz.component.shutdown.graceful.tomcat;

import lombok.extern.slf4j.Slf4j;
import org.apache.catalina.connector.Connector;
import org.apache.tomcat.util.threads.ThreadPoolExecutor;
import org.springframework.boot.web.embedded.tomcat.TomcatConnectorCustomizer;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextClosedEvent;

import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;

/**
 * @ClassName TomcatShutdown
 * @Description
 * @Author wzq
 * @Date 2023/9/25 13:36
 * @Version 1.0
 */
@Slf4j
public class TomcatShutdown implements TomcatConnectorCustomizer, ApplicationListener<ContextClosedEvent> {

    private volatile boolean shutdown = false;
    private int timeout;
    private volatile Connector connector;

    public TomcatShutdown(int timeout) {
        this.timeout = timeout;
    }

    @Override
    public void customize(Connector connector) {
        this.connector = connector;
    }

    @Override
    public void onApplicationEvent(ContextClosedEvent event) {
        try {
            if (!this.shutdown) {
                this.shutdown = true;
                long beginTime = System.currentTimeMillis();
                log.info("tomcat start beginTime [{}]", beginTime);
                this.connector.pause();
                Executor executor = this.connector.getProtocolHandler().getExecutor();
                if (executor instanceof ThreadPoolExecutor) {
                    try {
                        ThreadPoolExecutor threadPoolExecutor = (ThreadPoolExecutor) executor;
                        threadPoolExecutor.shutdown();
                        if (!threadPoolExecutor.awaitTermination(timeout, TimeUnit.SECONDS)) {
                            threadPoolExecutor.shutdownNow();
                        }
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                    }
                }
                Long endTime = System.currentTimeMillis() - beginTime;
                log.info("tomcat shutdown time [{}](ms)",endTime);
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        }

    }
}
